#include "Diffie.hpp"


using namespace std;

const static int logging = 1;
unsigned int private_key = computeRandomNumber();

// nodes that are in progress 
vector<PROGTABLE> progressTable;



void initiateDiffie () {

	WLANAddr destAddr [ WLAN_ADDR_LEN ];
	PCryptoPacket pcp = NULL;

	pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
	if ( pcp != NULL )
	{
		for (int i = 0; i < WLAN_ADDR_LEN; i++ )
		{
			destAddr->data[i] = 0xFF;
		}

		/* If they don't already exist, establish secure connection */
		if (checkSecureTable(destAddr) == -1)
		{
			pcp->id = DIFFIE_INIT | (DIFFIEHELLMAN << 4);
			sendDiffieCryptoPacket( pcp, destAddr );

			printf ("Sent DIFFIE Initialization\n");

		}
	}

}


int checkProgressTable (WLANAddr * addr) {
	for (int i = 0; i < progressTable.size(); i++) {
		if (progressTable.at(i).address.data[0] == addr->data[0] &&
  			progressTable.at(i).address.data[1] == addr->data[1] &&		
			progressTable.at(i).address.data[2] == addr->data[2] &&		
			progressTable.at(i).address.data[3] == addr->data[3] &&
			progressTable.at(i).address.data[4] == addr->data[4] &&
			progressTable.at(i).address.data[5] == addr->data[5])
				return i;
	}
	return -1;
}

unsigned int charToInt (char * in) {
	return atoi(in);
}

unsigned int keyMaker (unsigned int base, unsigned int exp, unsigned int mod ){
		
	return modularExponentiation(base, exp, mod);

}

void addToSecureTable (WLANAddr * srcAddr, unsigned int key, double time) {

    char asciiAddr[32];

	char * dest = srcAddr->wlan2asc();
	strcpy(asciiAddr, dest);

	if (checkSecureTable(srcAddr) == -1)
	{
		printf ("[ Add to secure table: %s ]\n", asciiAddr);
		tableEntry temp = {*srcAddr, DIFFIEHELLMAN, key, time};
		pushToSecureTable (temp);
		return;
	}

	printf ("[ Address %s already added ]\n", asciiAddr);
}

void popFromProgressTable (int id) {
		// Swap and pop to remove from Progress stack
		PROGTABLE temp = {NULL,0,0,0,0};
		temp = progressTable.at(progressTable.size()-1);
		progressTable.at(progressTable.size()-1) = progressTable.at(id);
		progressTable.at(id) = temp;
		progressTable.pop_back();
}

void handleDH ( PCryptoPacket pCryptoPacket, WLANHeader * srcHeader ) { 

	WLANAddr srcAddr = srcHeader->srcAddr;
	PCryptoPacket pcp = NULL;
	
	int id = checkSecureTable(&srcAddr);

	if (id != -1)
		popFromSecureTable(id);

	char * buff = pCryptoPacket->data;
	unsigned short type = pCryptoPacket->id & 0x0F;

	// Existing Entry [DIFFIE_SEND]
	if ( type == DIFFIE_SEND ) {
		printf ("\n[ Existing Node ]: id=%d\n",id);
	
		// decrypt message and have fun

	
		return;
	}
	id = checkProgressTable (&srcAddr);

	// In Progress Entry [DIFFIE_PUBLIC_KEY] - if they're in progress, get public key, calculate shared key, add/remove
	if ( type == DIFFIE_PUBLICKEY  && id != -1) {
		printf ("\n[ Progress Node ]: id=%d\n",id);

		// Extract Public Key
		unsigned int public_key;
		memcpy( &public_key, pCryptoPacket->data, sizeof(unsigned int)); 
		printf ("received public_key=%u\n",public_key); 	
	
		// calculate shared key + add
		struct timeval endTime;
		gettimeofday( &endTime, NULL );		
		double timeTaken ((double)endTime.tv_usec - (double)progressTable.at(id).time);
		timeTaken /= 1000000;

		unsigned int shared_key = keyMaker ( public_key, private_key, progressTable.at(id).p );
		printf ("shared_key=%u\n",shared_key);	
		addToSecureTable (&srcAddr, shared_key, timeTaken);

		popFromProgressTable(id);

		printf ("RUNNING_TIME:%f\n", timeTaken);

	}

	// New Entry [DIFFIE_PG_PUBLICKEY] - retreive p, g, public key from message, then send our public key
	else if ( type == DIFFIE_PG_PUBLICKEY ) { 
			printf ("\n[ Connection Requested ]\n");

			// Strip p, g, public_key values
			unsigned int p = 1;
			unsigned int g = 1;
			unsigned int public_key = 1;


			// generate own public_key
			memcpy( &p, pCryptoPacket->data, sizeof(unsigned int)); 
			memcpy( &g, pCryptoPacket->data + (sizeof(unsigned int)), sizeof(unsigned int)); 
			memcpy( &public_key, pCryptoPacket->data + (sizeof(unsigned int)) * 2, sizeof(unsigned int)); 



			printf ("p=%u\n",p);
			printf ("g=%u\n",g);
			printf ("received public_key=%u\n",public_key); 

			unsigned int shared_key = keyMaker (public_key, private_key, p);
			printf ("shared_key=%u\n",shared_key);
			addToSecureTable (&srcAddr, shared_key, -1);

			

			// generate message then send
			pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
			if ( pcp != NULL )
			{

				pcp->id = DIFFIE_PUBLICKEY | (DIFFIEHELLMAN << 4);
				unsigned int own_public_key = keyMaker(g, private_key, p);
				printf ("own public_key=%u\n",own_public_key); 
				printf ("own private_key=%u\n",private_key); 

				// add the data (this node's public key)
				memcpy( pcp->data, &own_public_key, sizeof(unsigned int) );

				// send it
				printf ("*Sending Public Key*\n");
				sendDiffieCryptoPacket( pcp, &srcAddr );

				// calculate shared key + add to secure table


			}
			else
				printf ("PCP failed malloc");




		} 
	// Initiator [DIFFIE_INIT] - calculates p, g, public key, then sends
	else if ( type == DIFFIE_INIT ) {
		
		printf ("\n[ Initiating request ]\n");
		unsigned int p = computeRandomNumber();
		unsigned int g = computeRandomNumber();
		unsigned int public_key = keyMaker(g, private_key, p);
		printf ("p=%u\n",p);
		printf ("g=%u\n",g);
		printf ("own public_key=%u\n",public_key);


		struct timeval startTime;
		gettimeofday( &startTime, NULL );		
		PROGTABLE temp = {srcAddr, p, g, 0, (double)startTime.tv_usec};

		progressTable.push_back(temp);

		// create message

		pcp = (CryptoPacket *)malloc(sizeof(CryptoPacket));
		if ( pcp != NULL )
		{

			pcp->id = DIFFIE_PG_PUBLICKEY | (DIFFIEHELLMAN << 4);

			// add the data (p, g, public key)
			memcpy( pcp->data, &p, sizeof(unsigned int) );
			memcpy( pcp->data + sizeof(unsigned int), &g, sizeof(unsigned int) );
			memcpy( pcp->data + sizeof(unsigned int)*2, &public_key, sizeof(unsigned int) );
			
			// send it
			printf ("*Sending PGPK*\n");
			sendDiffieCryptoPacket( pcp, &srcAddr);
		}

	}


}

void sendDiffieCryptoPacket( PCryptoPacket pCryptoPacket, WLANAddr * destAddr)
{

    char asciiAddr[32];

	char * dest = destAddr->wlan2asc();
	strcpy(asciiAddr, dest);
	printf ("Send to: %s\n", asciiAddr);
	


	void * pBuffer = malloc( sizeof(CryptoPacket) );
	if ( pBuffer != NULL )
	{
		memcpy( (char*)pBuffer, pCryptoPacket, sizeof(CryptoPacket) );
		MySendData( destAddr->data, pBuffer );
	}

	if ( pCryptoPacket != NULL )
	{
		free( pCryptoPacket );
	}

}





